import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def plot_robot_arm(DH_params):
    # 提取DH参数
    d = DH_params[:, 3]
    a = DH_params[:, 2]
    alpha = np.deg2rad(DH_params[:, 1])  # 将角度转换为弧度
    theta = np.deg2rad(DH_params[:, 4])  # 将角度转换为弧度

    # 初始化变换矩阵
    T = np.eye(4)

    # 定义一个存储每个关节位置的列表
    positions = [np.zeros(3)]  # 初始化位置列表，第一个元素是原点

    # 计算每个关节的位置
    for i in range(len(DH_params)):
        A_i = np.array([
            [np.cos(theta[i]), -np.sin(theta[i]) * np.cos(alpha[i]), np.sin(theta[i]) * np.sin(alpha[i]), a[i] * np.cos(theta[i])],
            [np.sin(theta[i]), np.cos(theta[i]) * np.cos(alpha[i]), -np.cos(theta[i]) * np.sin(alpha[i]), a[i] * np.sin(theta[i])],
            [0, np.sin(alpha[i]), np.cos(alpha[i]), d[i]],
            [0, 0, 0, 1]
        ])
        T = T @ A_i
        positions.append(T[:3, 3])

    # 绘制机械臂
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    # 绘制每个关节到下一个关节的线段
    for i in range(len(positions) - 1):
        ax.plot([positions[i][0], positions[i+1][0]], [positions[i][2], positions[i+1][2]], [positions[i][1], positions[i+1][1]], 'bo-', linewidth=4)

    ax.set_xlabel('X (mm)')
    ax.set_ylabel('Z (mm)')
    ax.set_zlabel('Y (mm)')
    plt.title('六自由度机械臂的零位状态')
    ax.set_xlim([-500, 2000])
    ax.set_ylim([-500, 3000])
    ax.set_zlim([-1000, 1000])
    ax.view_init(elev=0, azim=90)  # 设置视角

    plt.grid(True)
    plt.show()

if __name__ == '__main__':
    DH_params = np.array([
        [1, -45, 0, 560, 10],
        [2, 45, 0, 450, 0],
        [3, 45, 0, 410, 10],
        [4, -45, 0, 380, -20],
        [5, 45, 0, 360, 0],
        [6, 0, 0, 210, 40]
    ])
    plot_robot_arm(DH_params)